import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, RandomSampler
import matplotlib.pyplot as plt
import numpy as np
import torchvision
from functools import partial
# 定义数据预处理
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(mean=(0.5,), std=(0.5,))  # 将图像归一化到[-1, 1]范围内
])

# 下载并加载训练数据
train_dataset = datasets.MNIST(root='./data', train=True, transform=transform, download=False)
sample_index = [i*20 for i in range(3000)] #假设取前500个训练数据
X_train = []
y_train = []
for i in sample_index:
    X = train_dataset[i][0]
    X_train.append(X)
    y = train_dataset[i][1]
    y_train.append(y)

sampled_train_data = [(X, y) for X, y in zip(X_train, y_train)] #包装为数据对
trainDataLoader = torch.utils.data.DataLoader(sampled_train_data, batch_size=16, shuffle=True)

train_loader = DataLoader(dataset=sampled_train_data, batch_size=128, shuffle=True)

class Generator(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(Generator, self).__init__()
        self.main = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, output_size),
            nn.Tanh()
        )

    def forward(self, x):
        return self.main(x)


class Discriminator(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(Discriminator, self).__init__()
        self.main = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.LeakyReLU(0.2),
            nn.Linear(hidden_size, hidden_size),
            nn.LeakyReLU(0.2),
            nn.Linear(hidden_size, output_size),
            nn.Sigmoid()
        )

    def forward(self, x):
        return self.main(x)

# 定义模型参数
input_size = 100  # 噪声向量的维度
hidden_size = 256
image_size = 28 * 28  # MNIST图像的维度

# 创建生成器和判别器实例
G = Generator(input_size, hidden_size, image_size)
D = Discriminator(image_size, hidden_size, 1)

criterion = nn.BCELoss()
lr = 0.0005

# 创建生成器和判别器的优化器
optimizer_G = optim.Adam(G.parameters(), lr=lr)
optimizer_D = optim.Adam(D.parameters(), lr=lr)

num_epochs = 1000

for epoch in range(num_epochs):
    for i, (images, _) in enumerate(train_loader):
        batch_size = images.size(0)
        images = images.view(batch_size, -1)

        # 创建标签
        real_labels = torch.ones(batch_size, 1)
        fake_labels = torch.zeros(batch_size, 1)

        # 训练判别器
        outputs = D(images)
        d_loss_real = criterion(outputs, real_labels)
        real_score = outputs

        z = torch.randn(batch_size, input_size)
        fake_images = G(z)
        outputs = D(fake_images.detach())
        d_loss_fake = criterion(outputs, fake_labels)
        fake_score = outputs

        d_loss = d_loss_real + d_loss_fake
        optimizer_D.zero_grad()
        d_loss.backward()
        optimizer_D.step()

        # 训练生成器
        z = torch.randn(batch_size, input_size)
        fake_images = G(z)
        outputs = D(fake_images)
        g_loss = criterion(outputs, real_labels)

        optimizer_G.zero_grad()
        g_loss.backward()
        optimizer_G.step()

    print(f'Epoch [{epoch+1}/{num_epochs}], d_loss: {d_loss.item():.4f}, g_loss: {g_loss.item():.4f}, '
          f'D(x): {real_score.mean().item():.4f}, D(G(z)): {fake_score.mean().item():.4f}')


def save_models(G, D, G_path='generator.pth', D_path='discriminator.pth'):
    """
    保存生成器和判别器模型

    参数:
        G: 训练好的生成器模型
        D: 训练好的判别器模型
        G_path: 生成器模型保存路径
        D_path: 判别器模型保存路径
    """
    # 保存生成器
    torch.save(G.state_dict(), G_path)
    # 保存判别器
    torch.save(D.state_dict(), D_path)
    print(f"模型已保存到 {G_path} 和 {D_path}")


# 保存模型
save_models(G, D)

# 生成一些新图像
z = torch.randn(64, input_size)
fake_images = G(z)
fake_images = fake_images.view(fake_images.size(0), 1, 28, 28)

# 可视化生成的图像
grid = torchvision.utils.make_grid(fake_images, nrow=8, normalize=True)
plt.imshow(grid.permute(1, 2, 0).detach().numpy())
plt.title('Generated Images')
plt.show()


